import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QPushButton, QFileDialog, QRadioButton, QButtonGroup,
                             QMessageBox, QProgressBar, QComboBox, QCheckBox)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QThread, pyqtSignal
import time

class MattingWorker(QThread):
    """处理抠图任务的工作线程"""
    progress_updated = pyqtSignal(int)
    task_completed = pyqtSignal(str, str)
    task_failed = pyqtSignal(str)

    def __init__(self, matting_processor, input_path, output_path, method, model_name=None, use_gpu=False):
        super().__init__()
        self.matting_processor = matting_processor
        self.input_path = input_path
        self.output_path = output_path
        self.method = method
        self.model_name = model_name
        self.use_gpu = use_gpu

    def run(self):
        try:
            start_time = time.time()
            if self.method == "u2net":
                result_path = self.matting_processor.process_image_u2net(
                    self.input_path, self.output_path, self.model_name, self.use_gpu)# 还未添加进度条，记得加self.progress_update.emit
                self.task_completed.emit(self.input_path, result_path)

            else:

                if os.path.isfile(self.input_path):
                    # 处理单个文件时更新进度
                    self.progress_updated.emit(0)
                    result_path = self.matting_processor.process_image_grabcut(
                        self.input_path, self.output_path)
                    self.task_completed.emit(self.input_path, result_path)
                    self.progress_updated.emit(100)
                else:
                    files = [f for f in os.listdir(self.input_path)
                             if os.path.isfile(os.path.join(self.input_path, f))]
                    total_files = len(files)
                    if total_files == 0:
                        raise ValueError("所选文件夹不包含任何文件")

                    for i, filename in enumerate(files):
                        file_path = os.path.join(self.input_path, filename)
                        result_path = self.matting_processor.process_image_grabcut(
                            file_path, self.output_path)
                        self.task_completed.emit(file_path, result_path)
                        self.progress_updated.emit(int((i + 1) / total_files * 100))
            print(f"本次处理用时{time.time()-start_time:.2f}秒")
        except Exception as e:
            print(e)
            self.task_failed.emit(str(e))


class MattingGUI(QMainWindow):
    """图像抠图工具的主窗口"""

    def __init__(self, matting_processor):
        super().__init__()
        self.output_path = None
        self.input_path = None
        self.matting_processor = matting_processor
        self.init_ui()

    def init_ui(self):
        # 设置窗口基本属性
        self.setWindowTitle('智能图像抠图工具')
        self.setGeometry(100, 100, 1000, 600)

        # 创建中心部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)

        # 左侧控制面板
        control_panel = QWidget()
        control_panel.setMinimumWidth(300)  # 设置最小宽度
        control_layout = QVBoxLayout(control_panel)

        # 输入类型选择
        input_type_label = QLabel('输入类型:')
        control_layout.addWidget(input_type_label)

        self.file_radio = QRadioButton('单个文件')
        self.folder_radio = QRadioButton('文件夹')
        self.file_radio.setChecked(True)

        self.input_type_group = QButtonGroup(self)
        self.input_type_group.addButton(self.file_radio, 0)
        self.input_type_group.addButton(self.folder_radio, 1)

        control_layout.addWidget(self.file_radio)
        control_layout.addWidget(self.folder_radio)

        # 输入路径选择
        input_path_layout = QHBoxLayout()
        self.input_path_edit = QLabel('未选择文件/文件夹')
        self.input_path_edit.setWordWrap(True)
        self.browse_button = QPushButton('浏览...')
        self.browse_button.clicked.connect(self.browse_input)

        input_path_layout.addWidget(self.input_path_edit)
        input_path_layout.addWidget(self.browse_button)
        control_layout.addLayout(input_path_layout)

        # 输出路径选择
        output_path_layout = QHBoxLayout()
        self.output_path_edit = QLabel('未设置')
        self.output_path_edit.setWordWrap(True)
        self.output_browse_button = QPushButton('浏览...')
        self.output_browse_button.clicked.connect(self.browse_output)

        output_path_layout.addWidget(self.output_path_edit)
        output_path_layout.addWidget(self.output_browse_button)
        control_layout.addLayout(output_path_layout)

        # 处理方法选择
        method_label = QLabel('抠图方法:')
        control_layout.addWidget(method_label)

        self.method_combo = QComboBox()
        self.method_combo.addItems(['OpenCV GrabCut', 'U2Net'])
        self.method_combo.currentIndexChanged.connect(self.on_method_changed)
        control_layout.addWidget(self.method_combo)

        # U2Net模型选择 (仅在选择U2Net方法时可见)
        self.model_label = QLabel('U2Net模型:')
        self.model_combo = QComboBox()
        self.model_combo.addItems(['u2net', 'u2netp'])
        control_layout.addWidget(self.model_label)
        control_layout.addWidget(self.model_combo)

        # GPU选择 (仅在选择U2Net方法时可见)
        self.gpu_checkbox = QCheckBox('使用GPU加速 (需要CUDA)')
        control_layout.addWidget(self.gpu_checkbox)

        # 初始化控件可见性
        self.on_method_changed(0)

        # 处理按钮
        self.process_button = QPushButton('开始抠图')
        self.process_button.clicked.connect(self.process)
        self.process_button.setMinimumHeight(40)
        control_layout.addWidget(self.process_button)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        control_layout.addWidget(self.progress_bar)

        control_layout.addStretch()

        # 右侧图像显示区域
        display_panel = QWidget()
        display_layout = QHBoxLayout(display_panel)

        # 原始图像显示
        self.original_label = QLabel('原始图像')
        self.original_label.setAlignment(Qt.AlignCenter)
        self.original_label.setMinimumSize(300, 400)
        self.original_label.setStyleSheet("border: 1px solid #cccccc;")

        # 结果图像显示
        self.result_label = QLabel('结果图像')
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setMinimumSize(300, 400)
        self.result_label.setStyleSheet("border: 1px solid #cccccc;")

        display_layout.addWidget(self.original_label)
        display_layout.addWidget(self.result_label)

        # 添加到主布局
        main_layout.addWidget(control_panel)
        main_layout.addWidget(display_panel)

        # 设置伸缩因子，让两个面板可以根据窗口大小调整
        main_layout.setStretch(0, 1)
        main_layout.setStretch(1, 3)

        # 初始化变量
        self.input_path = None
        self.output_path = None
        self.worker = None

    def on_method_changed(self, index):
        """处理方法变更事件"""
        use_u2net = (index == 1)
        self.model_label.setVisible(use_u2net)
        self.model_combo.setVisible(use_u2net)
        self.gpu_checkbox.setVisible(use_u2net)

    def browse_input(self):
        """浏览并选择输入文件或文件夹"""
        if self.file_radio.isChecked():
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择图像文件", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
            )
            if file_path:
                self.input_path = file_path
                self.input_path_edit.setText(os.path.basename(file_path))
                self.update_original_image(file_path)
                self.auto_set_output_path()
        else:
            folder_path = QFileDialog.getExistingDirectory(
                self, "选择文件夹", ""
            )
            if folder_path:
                self.input_path = folder_path
                self.input_path_edit.setText(os.path.basename(folder_path))
                self.auto_set_output_path()

    def browse_output(self):
        """浏览并选择输出文件夹"""
        folder_path = QFileDialog.getExistingDirectory(
            self, "选择输出文件夹", ""
        )
        if folder_path:
            self.output_path = folder_path
            self.output_path_edit.setText(os.path.basename(folder_path))

    def auto_set_output_path(self):
        """根据输入路径自动设置输出路径"""
        if self.input_path:
            if os.path.isfile(self.input_path):
                self.output_path = os.path.dirname(self.input_path)
            else:
                self.output_path = os.path.join(self.input_path, "matting_results")
                if not os.path.exists(self.output_path):
                    os.makedirs(self.output_path)
            self.output_path_edit.setText(os.path.basename(self.output_path))

    def update_original_image(self, image_path):
        """更新原始图像显示"""
        pixmap = QPixmap(image_path)
        scaled_pixmap = pixmap.scaled(
            self.original_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.original_label.setPixmap(scaled_pixmap)

    def update_result_image(self, image_path):
        """更新结果图像显示"""
        pixmap = QPixmap(image_path)
        scaled_pixmap = pixmap.scaled(
            self.result_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        self.result_label.setPixmap(scaled_pixmap)

    def process(self):
        """开始处理图像抠图"""
        if not self.input_path:
            QMessageBox.warning(self, "警告", "请先选择输入文件或文件夹")
            return

        if not self.output_path:
            QMessageBox.warning(self, "警告", "请先设置输出路径")
            return

        # 获取处理方法
        method_index = self.method_combo.currentIndex()
        method = 'u2net' if method_index == 1 else 'grabcut'

        # 获取U2Net相关参数（仅在使用U2Net时需要）
        model_name = self.model_combo.currentText() if method == 'u2net' else None
        use_gpu = self.gpu_checkbox.isChecked() if method == 'u2net' else False

        # 禁用处理按钮
        self.process_button.setEnabled(False)
        self.progress_bar.setValue(0)

        # 创建工作线程
        self.worker = MattingWorker(self.matting_processor, self.input_path, self.output_path,
                                    method, model_name, use_gpu)
        self.worker.progress_updated.connect(self.update_progress)
        self.worker.task_completed.connect(self.on_task_completed)
        self.worker.task_failed.connect(self.on_task_failed)
        self.worker.finished.connect(self.on_worker_finished)
        self.worker.start()

    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)

    def on_task_completed(self, input_path, result_path):
        """任务完成时的回调"""
        # 如果是处理单个文件，更新结果图像显示
        if os.path.isfile(self.input_path):
            self.update_result_image(result_path)
        print(f"处理完成: {input_path} -> {result_path}")

    def on_task_failed(self, error_msg):
        """任务失败时的回调"""
        QMessageBox.critical(self, "错误", f"处理失败: {error_msg}")

    def on_worker_finished(self):
        """工作线程完成时的回调"""
        self.process_button.setEnabled(True)
        if os.path.isdir(self.input_path):
            QMessageBox.information(self, "完成", "批量处理已完成")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    # 这里需要一个matting_processor实例，由于没有具体实现，用None代替
    matting_processor = None
    gui = MattingGUI(matting_processor)
    gui.show()
    sys.exit(app.exec_())